import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
import lightgbm as lgb
from sklearn.metrics import mean_squared_error

# 加载数据
train = pd.read_csv('训练集.csv')
test = pd.read_csv('测试集.csv')

# 特征工程
train['便利设施数量']=train['便利设施'].apply(lambda x:len(x.lstrip('{').rstrip('}').split(',')))
test['便利设施数量']=test['便利设施'].apply(lambda x:len(x.lstrip('{').rstrip('}').split(',')))

no_features = ['数据ID', '价格','便利设施']

# 类别编码
data = pd.concat([train, test], axis=0)
for col in train.select_dtypes(include=['object']).columns:
    if col not in no_features:
        lb = LabelEncoder()
        lb.fit(data[col].astype(str))
        train[col] = lb.transform(train[col].astype(str))
        test[col] = lb.transform(test[col].astype(str))

features = [col for col in train.columns if col not in no_features]
X = train[features]
y = train['价格']

# 快速测试当前模型性能
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

params = {
    'learning_rate': 0.1,
    'boosting_type': 'gbdt',
    'objective': 'regression',
    'metric': 'mae',
    'feature_fraction': 0.6,
    'bagging_fraction': 0.8,
    'bagging_freq': 5,
    'num_leaves': 1000,
    'verbose': -1,
    'max_depth': -1,
    'seed': 2019,
}

model = lgb.LGBMRegressor(**params, n_estimators=1000)
model.fit(X_train, y_train, eval_set=[(X_val, y_val)], eval_metric='rmse', callbacks=[lgb.log_evaluation(100), lgb.early_stopping(50)])
y_pred = model.predict(X_val)
score = mean_squared_error(y_val, y_pred)
print(f'当前MSE分数: {score:.2f}')